{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Overview\n",
    "\n",
    "This a notebook that inspects the results of a WarpX simulation and a Warp simulation at the same time. Typically used to compare the results of both codes for the same set of parameters.\n",
    "\n",
    "# Instructions\n",
    "\n",
    "- Enter the paths for the WarpX and the Warp simulations\n",
    "- Execute the cells below one by one, by selecting them with your mouse and typing `Shift + Enter`"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": false
   },
   "outputs": [],
   "source": [
    "# Import statements\n",
    "import glob\n",
    "\n",
    "import yt\n",
    "\n",
    "yt.funcs.mylog.setLevel(50)\n",
    "import matplotlib.pyplot as plt\n",
    "from ipywidgets import IntSlider, RadioButtons, interact\n",
    "\n",
    "%matplotlib"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": false
   },
   "outputs": [],
   "source": [
    "# Define path way for WarpX results and find iterations\n",
    "path_warpx = \"/Users/mthevenet/warpX_directory/res_warpx/valid_pwfa/\"\n",
    "file_list_warpx = glob.glob(path_warpx + \"plt?????\")\n",
    "iterations_warpx = [\n",
    "    int(file_name[len(file_name) - 5 :]) for file_name in file_list_warpx\n",
    "]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": false
   },
   "outputs": [],
   "source": [
    "# Define path way for Warp results and find iterations\n",
    "path_warp = \"/Users/mthevenet/warp_results/valid_pwfa/\"\n",
    "file_list_warp = glob.glob(path_warp + \"diags/hdf5/data????????.h5\")\n",
    "iterations_warp = [\n",
    "    int(file_name[len(file_name) - 11 : len(file_name) - 3])\n",
    "    for file_name in file_list_warp\n",
    "]"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Functions to plot the fields\n",
    "\n",
    "Note that the data are loaded with yt and plotted using matplotlib"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": false
   },
   "outputs": [],
   "source": [
    "def plot_field(iteration, field, slicing_direction=\"y\"):\n",
    "    # First dataset\n",
    "    ds = yt.load(path_warpx + \"/plt%05d/\" % iteration)\n",
    "    all_data_level_0 = ds.covering_grid(\n",
    "        level=0, left_edge=ds.domain_left_edge, dims=ds.domain_dimensions\n",
    "    )\n",
    "\n",
    "    # Second dataset\n",
    "    #     ts = OpenPMDTimeSeries(path_warp + 'diags/hdf5/')\n",
    "    ds2 = yt.load(path_warp + \"diags/hdf5/data%08d.h5\" % iteration)\n",
    "    all_data_level_02 = ds2.covering_grid(\n",
    "        level=0, left_edge=ds2.domain_left_edge, dims=ds2.domain_dimensions\n",
    "    )\n",
    "\n",
    "    # first dataset loaded via yt\n",
    "    left_edge = ds.domain_left_edge.convert_to_mks() * 1e6\n",
    "    right_edge = ds.domain_right_edge.convert_to_mks() * 1e6\n",
    "    if slicing_direction == \"x\":\n",
    "        n = int(ds.domain_dimensions[0] // 2)\n",
    "        data2d = all_data_level_0[field][n, :, :]\n",
    "        extent = [left_edge[2], right_edge[2], left_edge[1], right_edge[1]]\n",
    "    elif slicing_direction == \"y\":\n",
    "        n = int(ds.domain_dimensions[1] // 2)\n",
    "        data2d = all_data_level_0[field][:, n, :]\n",
    "        extent = [left_edge[2], right_edge[2], left_edge[0], right_edge[0]]\n",
    "    elif slicing_direction == \"z\":\n",
    "        n = int(ds.domain_dimensions[2] // 2)\n",
    "        data2d = all_data_level_0[field][:, :, n]\n",
    "        extent = [left_edge[1], right_edge[1], left_edge[0], right_edge[0]]\n",
    "\n",
    "    # second dataset loaded via yt\n",
    "    left_edge = ds2.domain_left_edge.convert_to_mks() * 1e6\n",
    "    right_edge = ds2.domain_right_edge.convert_to_mks() * 1e6\n",
    "    field_reformat = field[0].upper() + \"_\" + field[1]\n",
    "    if slicing_direction == \"x\":\n",
    "        n = int(ds2.domain_dimensions[0] // 2)\n",
    "        data2d2 = all_data_level_02[field_reformat][n, :, :]\n",
    "        extent2 = [left_edge[2], right_edge[2], left_edge[1], right_edge[1]]\n",
    "    elif slicing_direction == \"y\":\n",
    "        n = int(ds2.domain_dimensions[1] // 2)\n",
    "        data2d2 = all_data_level_02[field_reformat][:, n, :]\n",
    "        extent2 = [left_edge[2], right_edge[2], left_edge[0], right_edge[0]]\n",
    "    elif slicing_direction == \"z\":\n",
    "        n = int(ds2.domain_dimensions[2] // 2)\n",
    "        data2d2 = all_data_level_02[field_reformat][:, :, n]\n",
    "        extent2 = [left_edge[1], right_edge[1], left_edge[0], right_edge[0]]\n",
    "\n",
    "    if slicing_direction == \"x\":\n",
    "        yaxis_label = \"y\"\n",
    "    if slicing_direction == \"y\":\n",
    "        yaxis_label = \"x\"\n",
    "    if slicing_direction == \"z\":\n",
    "        yaxis_label = \"y\"\n",
    "\n",
    "    xlim = [extent[0], extent[1]]\n",
    "    ylim = [extent[2], extent[3]]\n",
    "\n",
    "    plt.clf()\n",
    "    # first dataset\n",
    "    plt.subplot(2, 1, 1)\n",
    "    plt.title(\"WarpX %s at iteration %d\" % (field, iteration))\n",
    "    plt.imshow(\n",
    "        data2d,\n",
    "        interpolation=\"nearest\",\n",
    "        cmap=\"viridis\",\n",
    "        origin=\"lower\",\n",
    "        extent=extent,\n",
    "        aspect=\"auto\",\n",
    "    )\n",
    "    plt.ylabel(yaxis_label + \" (um)\")\n",
    "    vmin, vmax = plt.gci().get_clim()\n",
    "    plt.colorbar()\n",
    "    if \"xlim\" in locals():\n",
    "        plt.xlim(xlim)\n",
    "    if \"ylim\" in locals():\n",
    "        plt.ylim(ylim)\n",
    "    plt.xticks([])\n",
    "\n",
    "    # second dataset\n",
    "    plt.subplot(2, 1, 2)\n",
    "    plt.title(\"Warp %s at iteration %d\" % (field, iteration))\n",
    "    plt.imshow(\n",
    "        data2d2,\n",
    "        interpolation=\"nearest\",\n",
    "        cmap=\"viridis\",\n",
    "        origin=\"lower\",\n",
    "        extent=extent2,\n",
    "        aspect=\"auto\",\n",
    "        vmin=vmin,\n",
    "        vmax=vmax,\n",
    "    )\n",
    "    plt.colorbar()\n",
    "    if \"xlim\" in locals():\n",
    "        plt.xlim(xlim)\n",
    "    if \"ylim\" in locals():\n",
    "        plt.ylim(ylim)\n",
    "    plt.xlabel(\"z (um)\")\n",
    "    plt.ylabel(yaxis_label + \" (um)\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Interactive viewer"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": false
   },
   "outputs": [],
   "source": [
    "iterations = iterations_warp\n",
    "interact(\n",
    "    plot_field,\n",
    "    iteration=IntSlider(\n",
    "        min=min(iterations), max=max(iterations), step=iterations[1] - iterations[0]\n",
    "    ),\n",
    "    field=RadioButtons(options=[\"jx\", \"jy\", \"jz\", \"Ex\", \"Ey\", \"Ez\"], value=\"jx\"),\n",
    "    slicing_direction=RadioButtons(options=[\"x\", \"y\", \"z\"], value=\"x\"),\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "anaconda-cloud": {},
  "kernelspec": {
   "display_name": "Python 2",
   "language": "python",
   "name": "python2"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 2
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython2",
   "version": "2.7.13"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 1
}
